#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class ClipByValue {
 public:
  __aicore__ inline ClipByValue() {}
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR clip_value_min, GM_ADDR clip_value_max, GM_ADDR y,
                              uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t tilingKey) {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->value = (DTYPE_X)0;

    this->blockLength = blockLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;

    xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    minGm.SetGlobalBuffer((__gm__ DTYPE_X*)clip_value_min + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    maxGm.SetGlobalBuffer((__gm__ DTYPE_X*)clip_value_max + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    yGm.SetGlobalBuffer(
        (__gm__ DTYPE_X*)y + this->blockLength * GetBlockIdx(),
        this->blockLength);
  


    pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));

    pipe.InitBuffer(Maxn, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(Minn, this->tileLength * sizeof(DTYPE_X));

    Init_m(this->value);
  }

  template<typename T>
  __aicore__ inline void Init_m(T flag) {
      this->minn = minGm.GetValue(0);
      this->maxn = maxGm.GetValue(0);
      LocalTensor<T> minLocal = Minn.Get<T>();
      LocalTensor<T> maxLocal = Maxn.Get<T>();
      Duplicate(maxLocal, this->maxn, this->tileLength);
      Duplicate(minLocal, this->minn, this->tileLength);
  }
  template<>
  __aicore__ inline void Init_m(int32_t flag) {
      this->minn_32 = minGm.GetValue(0);
      this->maxn_32 = maxGm.GetValue(0);
      LocalTensor<int32_t> minLocal = Minn.Get<int32_t>();
      LocalTensor<int32_t> maxLocal = Maxn.Get<int32_t>();
      Duplicate(maxLocal, this->maxn_32, this->tileLength);
      Duplicate(minLocal, this->minn_32, this->tileLength);
  }

  __aicore__ inline void Process() {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++) {
      CopyIn(i);
      Compute(i, value);
      CopyOut(i);
    }
  }

 private:
  __aicore__ inline void CopyIn(int32_t progress) {
    LocalTensor<DTYPE_X> inLocal = inQueueIN.AllocTensor<DTYPE_X>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          DataCopy(inLocal[0], xGm[0], this->tileLength);
        } else {
          DataCopy(
              inLocal[0],
              xGm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
        }
      } else {
        DataCopy(inLocal[0], xGm[progress * this->tileLength],
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            inLocal[0],
            xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
      }

      else {
        DataCopy(inLocal[0], xGm[progress * (this->tileLength)],
                 (this->tileLength));
      }
    }

    inQueueIN.EnQue(inLocal);
  }
  template<typename T>
  __aicore__ inline void Compute(int32_t progress, T flag) {
    LocalTensor<T> inLocal = inQueueIN.DeQue<T>();
    LocalTensor<T> xLocal = inLocal;
    LocalTensor<T> minLocal = Minn.Get<T>();
    LocalTensor<T> maxLocal = Maxn.Get<T>();
    LocalTensor<T> outLocal = outQueueOUT.AllocTensor<T>();

    Max(outLocal, minLocal, xLocal, this->tileLength);
    Min(outLocal, maxLocal, outLocal, this->tileLength);

    outQueueOUT.EnQue<T>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }
  template <>
  __aicore__ inline void Compute(int32_t progress, int32_t flag) {
    LocalTensor<int32_t> inLocal = inQueueIN.DeQue<int32_t>();
    LocalTensor<int32_t> xLocal = inLocal;
    LocalTensor<int32_t> minLocal = Minn.Get<int32_t>();
    LocalTensor<int32_t> maxLocal = Maxn.Get<int32_t>();

    LocalTensor<int32_t> outLocal = outQueueOUT.AllocTensor<int32_t>();
    Max(outLocal, minLocal, xLocal, this->tileLength);
    Min(outLocal, maxLocal, outLocal, this->tileLength);

    outQueueOUT.EnQue<int32_t>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }

  __aicore__ inline void CopyOut(int32_t progress) {
    LocalTensor<DTYPE_X> outLocal = outQueueOUT.DeQue<DTYPE_X>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          DataCopy(yGm[0], outLocal, this->tileLength);
        } else {
          DataCopy(
              yGm[(progress - 1) * this->tileLength + this->lasttileLength],
              outLocal, this->tileLength);
        }
      } else {
        DataCopy(yGm[progress * this->tileLength], outLocal,
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            outLocal, (this->tileLength));
      }

      else {
        DataCopy(yGm[progress * (this->tileLength)], outLocal,
                 (this->tileLength));
      }
    }

    outQueueOUT.FreeTensor(outLocal);
  }

 private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
  GlobalTensor<DTYPE_X> xGm;
  GlobalTensor<DTYPE_X> minGm;
  GlobalTensor<DTYPE_X> maxGm;
  GlobalTensor<DTYPE_X> yGm;
  DTYPE_X value;

  DTYPE_X minn;
  DTYPE_X maxn;
  TBuf<> Maxn;
  TBuf<> Minn;

  int32_t minn_32;
  int32_t maxn_32;

  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
};


extern "C" __global__ __aicore__ void clip_by_value(GM_ADDR x, GM_ADDR clip_value_min, GM_ADDR clip_value_max, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    ClipByValue op;
    uint32_t tilingKey = 1;
    if (TILING_KEY_IS(1)) {
        tilingKey = 1;
    } else if (TILING_KEY_IS(2)) {
        tilingKey = 2;
    } else {
        tilingKey = 1;
    }

    op.Init(x, clip_value_min, clip_value_max, y,  tiling_data.blockLength,
            tiling_data.tileNum, tiling_data.tileLength,
            tiling_data.lasttileLength, 
            tilingKey);
    op.Process();
}